// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include "zencore/zencore.h"

#include "zencore/blake3.h"
#include "zencore/compositebuffer.h"
#include "zencore/intmath.h"

#include <limits>

namespace zen {

class Archive;

namespace detail {
	struct BufferHeader;
}

enum class OodleCompressor : uint8_t
{
	NotSet	  = 0,
	Selkie	  = 1,
	Mermaid	  = 2,
	Kraken	  = 3,
	Leviathan = 4,
};

enum class OodleCompressionLevel : int8_t
{
	HyperFast4 = -4,
	HyperFast3 = -3,
	HyperFast2 = -2,
	HyperFast1 = -1,
	None	   = 0,
	SuperFast  = 1,
	VeryFast   = 2,
	Fast	   = 3,
	Normal	   = 4,
	Optimal1   = 5,
	Optimal2   = 6,
	Optimal3   = 7,
	Optimal4   = 8,
};

/**
 * A compressed buffer stores compressed data in a self-contained format.
 *
 * A buffer is self-contained in the sense that it can be decompressed without external knowledge
 * of the compression format or the size of the raw data.
 */
class CompressedBuffer
{
public:
	/**
	 * Compress the buffer using the specified compressor and compression level.
	 *
	 * Data that does not compress will be return uncompressed, as if with level None.
	 *
	 * @note Using a level of None will return a buffer that references owned raw data.
	 *
	 * @param RawData            The raw data to be compressed.
	 * @param Compressor         The compressor to encode with. May use NotSet if level is None.
	 * @param CompressionLevel   The compression level to encode with.
	 * @param BlockSize          The power-of-two block size to encode raw data in. 0 is default.
	 * @return An owned compressed buffer, or null on error.
	 */
	[[nodiscard]] ZENCORE_API static CompressedBuffer Compress(const CompositeBuffer& RawData,
															   OodleCompressor		  Compressor	   = OodleCompressor::Mermaid,
															   OodleCompressionLevel  CompressionLevel = OodleCompressionLevel::VeryFast,
															   uint64_t				  BlockSize		   = 0);
	[[nodiscard]] ZENCORE_API static CompressedBuffer Compress(const SharedBuffer&	 RawData,
															   OodleCompressor		 Compressor		  = OodleCompressor::Mermaid,
															   OodleCompressionLevel CompressionLevel = OodleCompressionLevel::VeryFast,
															   uint64_t				 BlockSize		  = 0);

	/**
	 * Construct from a compressed buffer previously created by Compress().
	 *
	 * @return A compressed buffer, or null on error, such as an invalid format or corrupt header.
	 */
	[[nodiscard]] ZENCORE_API static CompressedBuffer FromCompressed(const CompositeBuffer& CompressedData,
																	 IoHash&				OutRawHash,
																	 uint64_t&				OutRawSize);
	[[nodiscard]] ZENCORE_API static CompressedBuffer FromCompressed(CompositeBuffer&& CompressedData,
																	 IoHash&		   OutRawHash,
																	 uint64_t&		   OutRawSize);
	[[nodiscard]] ZENCORE_API static CompressedBuffer FromCompressed(const SharedBuffer& CompressedData,
																	 IoHash&			 OutRawHash,
																	 uint64_t&			 OutRawSize);
	[[nodiscard]] ZENCORE_API static CompressedBuffer FromCompressed(SharedBuffer&& CompressedData,
																	 IoHash&		OutRawHash,
																	 uint64_t&		OutRawSize);
	[[nodiscard]] ZENCORE_API static CompressedBuffer FromCompressedNoValidate(IoBuffer&& CompressedData);
	[[nodiscard]] ZENCORE_API static CompressedBuffer FromCompressedNoValidate(CompositeBuffer&& CompressedData);
	[[nodiscard]] ZENCORE_API static bool	ValidateCompressedHeader(IoBuffer&& CompressedData, IoHash& OutRawHash, uint64_t& OutRawSize);
	[[nodiscard]] ZENCORE_API static bool	ValidateCompressedHeader(const IoBuffer& CompressedData,
																	 IoHash&		 OutRawHash,
																	 uint64_t&		 OutRawSize);
	[[nodiscard]] ZENCORE_API static size_t GetHeaderSizeForNoneEncoder();
	[[nodiscard]] ZENCORE_API static UniqueBuffer CreateHeaderForNoneEncoder(uint64_t RawSize, const BLAKE3& RawHash);

	/** Reset this to null. */
	inline void Reset() { CompressedData.Reset(); }

	/** Returns true if the compressed buffer is not null. */
	[[nodiscard]] inline explicit operator bool() const { return !IsNull(); }

	/** Returns true if the compressed buffer is null. */
	[[nodiscard]] inline bool IsNull() const { return CompressedData.IsNull(); }

	/** Returns true if the composite buffer is owned. */
	[[nodiscard]] inline bool IsOwned() const { return CompressedData.IsOwned(); }

	/** Returns a copy of the compressed buffer that owns its underlying memory. */
	[[nodiscard]] inline CompressedBuffer MakeOwned() const& { return FromCompressedNoValidate(CompressedData.MakeOwned()); }
	[[nodiscard]] inline CompressedBuffer MakeOwned() && { return FromCompressedNoValidate(std::move(CompressedData).MakeOwned()); }

	/** Returns a composite buffer containing the compressed data. May be null. May not be owned. */
	[[nodiscard]] inline const CompositeBuffer& GetCompressed() const& { return CompressedData; }
	[[nodiscard]] inline CompositeBuffer		GetCompressed() && { return std::move(CompressedData); }

	/** Returns the size of the compressed data. Zero if this is null. */
	[[nodiscard]] inline uint64_t GetCompressedSize() const { return CompressedData.GetSize(); }

	/** Returns the size of the raw data. Zero on error or if this is empty or null. */
	[[nodiscard]] ZENCORE_API uint64_t DecodeRawSize() const;

	/** Returns the hash of the raw data. Zero on error or if this is null. */
	[[nodiscard]] ZENCORE_API IoHash DecodeRawHash() const;

	/**
	 * Returns a block aligned range of a compressed buffer.
	 *
	 * This extracts a sub-range from the compressed buffer, if the buffer is block-compressed
	 * it will align start and end to end up on block boundaries.
	 *
	 * The resulting segments in the CompressedBuffer will are allocated and the data is copied
	 * from the source buffers.
	 *
	 * A new header will be allocated and generated.
	 *
	 * The RawHash field of the header will be zero as we do not calculate the raw hash for the sub-range
	 *
	 * @return A sub-range from the compressed buffer that encompasses RawOffset and RawSize
	 */
	[[nodiscard]] ZENCORE_API CompressedBuffer CopyRange(uint64_t RawOffset, uint64_t RawSize = ~uint64_t(0)) const;

	/**
	 * Returns a block aligned range of a compressed buffer.
	 *
	 * This extracts a sub-range from the compressed buffer, if the buffer is block-compressed
	 * it will align start and end to end up on block boundaries.
	 *
	 * The resulting segments in the CompressedBuffer will reference the source buffers so it won't
	 * allocate memory and copy data for the compressed data blocks.
	 *
	 * A new header will be allocated and generated.
	 *
	 * The RawHash field of the header will be zero as we do not calculate the raw hash for the sub-range
	 *
	 * @return A sub-range from the compressed buffer that encompasses RawOffset and RawSize
	 */
	[[nodiscard]] ZENCORE_API CompressedBuffer GetRange(uint64_t RawOffset, uint64_t RawSize = ~uint64_t(0)) const;

	/**
	 * Returns the compressor and compression level used by this buffer.
	 *
	 * The compressor and compression level may differ from those specified when creating the buffer
	 * because an incompressible buffer is stored with no compression. Parameters cannot be accessed
	 * if this is null or uses a method other than Oodle, in which case this returns false.
	 *
	 * @return True if parameters were written, otherwise false.
	 */
	[[nodiscard]] ZENCORE_API bool TryGetCompressParameters(OodleCompressor&	   OutCompressor,
															OodleCompressionLevel& OutCompressionLevel,
															uint64_t&			   OutBlockSize) const;

	/**
	 * Decompress into a memory view that is less or equal GetRawSize() bytes.
	 */
	[[nodiscard]] ZENCORE_API bool TryDecompressTo(MutableMemoryView RawView, uint64_t RawOffset = 0) const;

	/**
	 * Decompress into an owned buffer.
	 *
	 * @return An owned buffer containing the raw data, or null on error.
	 */
	[[nodiscard]] ZENCORE_API SharedBuffer Decompress(uint64_t RawOffset = 0, uint64_t RawSize = ~uint64_t(0)) const;

	/**
	 * Decompress into an owned composite buffer.
	 *
	 * @return An owned buffer containing the raw data, or null on error.
	 */
	[[nodiscard]] ZENCORE_API CompositeBuffer DecompressToComposite() const;

	/** A null compressed buffer. */
	static const CompressedBuffer Null;

private:
	CompositeBuffer CompressedData;
};

namespace detail {
	/** A reusable context for the compressed buffer decoder. */
	struct DecoderContext
	{
		/** The offset in the source at which the compressed buffer begins, otherwise MAX_uint64. */
		uint64_t HeaderOffset = MAX_uint64;
		/** The size of the header if known, otherwise 0. */
		uint64_t HeaderSize = 0;
		/** The CRC-32 from the header, otherwise 0. */
		uint32_t HeaderCrc32 = 0;
		/** Index of the block stored in RawBlock, otherwise MAX_uint32. */
		uint32_t RawBlockIndex = MAX_uint32;

		/** A buffer used to store the header when HeaderOffset is not MAX_uint64. */
		UniqueBuffer Header;
		/** A buffer used to store a raw block when a partial block read is requested. */
		UniqueBuffer RawBlock;
		/** A buffer used to store a compressed block when it was not in contiguous memory. */
		UniqueBuffer CompressedBlock;
	};
}  // namespace detail

/**
 * A type that stores the state needed to decompress a compressed buffer.
 *
 * The compressed buffer can be in memory or can be loaded from a seekable archive.
 *
 * The reader can be reused across multiple source buffers, which allows its temporary buffers to
 * be reused if they are the right size.
 *
 * It is only safe to use the reader from one thread at a time.
 *
 * @see CompressedBuffer
 */
class CompressedBufferReader
{
public:
	/** Construct a reader with no source. */
	CompressedBufferReader() = default;

	/** Construct a reader that will read from an archive as needed. */
	explicit CompressedBufferReader(Archive& Archive);

	/** Construct a reader from an in-memory compressed buffer. */
	explicit CompressedBufferReader(const CompressedBuffer& Buffer);

	/** Release any temporary buffers that have been allocated by the reader. */
	void ResetBuffers();

	/** Clears the reference to the source without releasing temporary buffers. */
	void ResetSource();

	void SetSource(Archive& Archive);
	void SetSource(const CompressedBuffer& Buffer);

	[[nodiscard]] inline bool HasSource() const { return SourceArchive || SourceBuffer; }

	/** Returns the size of the compressed data. Zero on error. */
	[[nodiscard]] uint64_t GetCompressedSize();

	/** Returns the size of the raw data. Zero on error. */
	[[nodiscard]] uint64_t GetRawSize();

	/** Returns the hash of the raw data. Zero on error. */
	[[nodiscard]] IoHash GetRawHash();

	/**
	 * Returns the compressor and compression level used by this buffer.
	 *
	 * The compressor and compression level may differ from those specified when creating the buffer
	 * because an incompressible buffer is stored with no compression. Parameters cannot be accessed
	 * if this is null or uses a method other than Oodle, in which case this returns false.
	 *
	 * @return True if parameters were written, otherwise false.
	 */
	[[nodiscard]] bool TryGetCompressParameters(OodleCompressor&	   OutCompressor,
												OodleCompressionLevel& OutCompressionLevel,
												uint64_t&			   OutBlockSize);

	/**
	 * Decompress into a memory view that is less than or equal to the available raw size.
	 *
	 * @param RawView     The view to write to. The size to read is equal to the view size.
	 * @param RawOffset   The offset into the raw data from which to decompress.
	 * @return True if the requested range was decompressed, otherwise false.
	 */
	[[nodiscard]] bool TryDecompressTo(MutableMemoryView RawView, uint64_t RawOffset = 0);

	/**
	 * Decompress into an owned buffer.
	 *
	 * RawOffset must be at most the raw buffer size. RawSize may be MAX_uint64 to read the whole
	 * buffer from RawOffset, and must otherwise fit within the bounds of the buffer.
	 *
	 * @param RawOffset   The offset into the raw data from which to decompress.
	 * @param RawSize     The size of the raw data to read from the offset.
	 * @return An owned buffer containing the raw data, or null on error.
	 */
	[[nodiscard]] SharedBuffer	  Decompress(uint64_t RawOffset = 0, uint64_t RawSize = MAX_uint64);
	[[nodiscard]] CompositeBuffer DecompressToComposite(uint64_t RawOffset = 0, uint64_t RawSize = MAX_uint64);

private:
	bool TryReadHeader(detail::BufferHeader& OutHeader, MemoryView& OutHeaderView);

	Archive*				SourceArchive = nullptr;
	const CompressedBuffer* SourceBuffer  = nullptr;
	detail::DecoderContext	Context;
};

void compress_forcelink();	// internal

}  // namespace zen
